{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Cartopy Examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from __future__ import (absolute_import, division, print_function, unicode_literals)\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.cm import get_cmap\n",
    "import cartopy.crs as crs\n",
    "import cartopy.feature as cfeature\n",
    "from netCDF4 import Dataset\n",
    "\n",
    "from wrf import to_np, getvar, smooth2d, get_cartopy, latlon_coords, CoordPair, GeoBounds, cartopy_xlim, cartopy_ylim\n",
    "\n",
    "# Open the NetCDF file\n",
    "ncfile = Dataset(\"/Users/ladwig/Documents/wrf_files/wrfout_d01_2016-10-07_00_00_00\")\n",
    "\n",
    "# Get sea level pressure and cloud top temperature\n",
    "slp = getvar(ncfile, \"slp\")\n",
    "ctt = getvar(ncfile, \"ctt\")\n",
    "\n",
    "# Smooth the SLP\n",
    "smooth_slp = smooth2d(slp, 3)\n",
    "\n",
    "# Extract the latitude and longitude coordinate arrays as regular numpy array instead of xarray.DataArray\n",
    "lats, lons = latlon_coords(slp)\n",
    "\n",
    "# Get the cartopy projection class\n",
    "cart_proj = get_cartopy(slp)\n",
    "\n",
    "# Create the figure\n",
    "fig = plt.figure(figsize=(8,8))\n",
    "ax = plt.axes(projection=cart_proj)\n",
    "\n",
    "# Download and create the states, land, and oceans using cartopy features\n",
    "states = cfeature.NaturalEarthFeature(category='cultural', scale='50m', facecolor='none',\n",
    "                             name='admin_1_states_provinces_shp')\n",
    "land = cfeature.NaturalEarthFeature(category='physical', name='land', scale='50m', \n",
    "                                    facecolor=cfeature.COLORS['land'])\n",
    "ocean = cfeature.NaturalEarthFeature(category='physical', name='ocean', scale='50m', \n",
    "                                     facecolor=cfeature.COLORS['water'])\n",
    "\n",
    "# Make the pressure contours.\n",
    "#contour_levels = [960, 965, 970, 975, 980, 990]\n",
    "#c1 = plt.contour(lons, lats, to_np(smooth_slp), levels=contour_levels, colors=\"white\", \n",
    "#            transform=crs.PlateCarree(), zorder=3, linewidths=1.0)\n",
    "\n",
    "# Add pressure contour labels\n",
    "#plt.clabel(c1, contour_levels, inline=True, fmt='%.0f', fontsize=7)\n",
    "\n",
    "# Create the filled cloud top temperature contours\n",
    "#contour_levels = [-80, -70, -60, -50, -40, -30, -20, -10, 0]\n",
    "contour_levels = np.arange(-80.0, 10.0, 10.0)\n",
    "plt.contourf(to_np(lons), to_np(lats), to_np(ctt), contour_levels, cmap=get_cmap(\"Greys\"),\n",
    "             transform=crs.PlateCarree(), zorder=2)\n",
    "\n",
    "plt.plot([-80,-77.8], [26.76,26.76], color=\"yellow\", marker=\"o\", transform=crs.PlateCarree(), zorder=3)\n",
    "\n",
    "# Create the label bar for cloud top temperature\n",
    "#cb2 = plt.colorbar(ax=ax, fraction=0.046, pad=0.04)\n",
    "cb2 = plt.colorbar(ax=ax, shrink=.90)\n",
    "\n",
    "# Draw the oceans, land, and states\n",
    "ax.add_feature(ocean)\n",
    "ax.add_feature(land)\n",
    "ax.add_feature(states, linewidth=.5, edgecolor=\"black\")\n",
    "\n",
    "# Crop the domain to the region around the hurricane\n",
    "hur_bounds = GeoBounds(CoordPair(lat=np.amin(to_np(lats)), lon=-85.0),\n",
    "                       CoordPair(lat=30.0, lon=-72.0))\n",
    "ax.set_xlim(cartopy_xlim(ctt, geobounds=hur_bounds))\n",
    "ax.set_ylim(cartopy_ylim(ctt, geobounds=hur_bounds))\n",
    "ax.gridlines(color=\"white\", linestyle=\"dotted\")\n",
    "\n",
    "# Add the title and show the image\n",
    "plt.title(\"Hurricane Matthew Cloud Top Temperature (degC) \")\n",
    "#plt.savefig(\"/Users/ladwig/Documents/workspace/wrf_python/doc/source/_static/images/matthew.png\",\n",
    "#           transparent=True, bbox_inches=\"tight\")\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import (absolute_import, division, print_function, unicode_literals)\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.cm import get_cmap\n",
    "import cartopy.crs as crs\n",
    "from cartopy.feature import NaturalEarthFeature\n",
    "from netCDF4 import Dataset\n",
    "\n",
    "from wrf import to_np, getvar, smooth2d, ll_to_xy, CoordPair, vertcross, getproj, get_proj_params, to_xy_coords, latlon_coords\n",
    "\n",
    "# Open the NetCDF file\n",
    "ncfile = Dataset(\"/Users/ladwig/Documents/wrf_files/wrfout_d01_2016-10-07_00_00_00\")\n",
    "\n",
    "# Extract model height, dbz, and wind speed\n",
    "z = getvar(ncfile, \"z\", timeidx=0)\n",
    "dbz = getvar(ncfile, \"dbz\", timeidx=0)\n",
    "wspd =  getvar(ncfile, \"uvmet_wspd_wdir\", units=\"kt\")[0,:]\n",
    "Z = 10**(dbz/10.)\n",
    "\n",
    "start_point = CoordPair(lat=26.75, lon=-80.0)\n",
    "end_point = CoordPair(lat=26.75, lon=-77.8)\n",
    "\n",
    "# Compute the vertical cross-section interpolation.  Also, include the lat/lon points along the cross-section.\n",
    "z_cross = vertcross(Z, z, wrfin=ncfile, start_point=start_point, end_point=end_point, latlon=True, meta=True)\n",
    "wspd_cross = vertcross(wspd, z, wrfin=ncfile, start_point=start_point, end_point=end_point, latlon=True, meta=True)\n",
    "dbz_cross = 10.0 * np.log10(z_cross)\n",
    "\n",
    "# Create the figure\n",
    "fig, axes = plt.subplots(nrows=2, ncols=1, figsize=(4,4))\n",
    "#ax = plt.axes([0.1,0.1,0.8,0.8])\n",
    "\n",
    "# Define the contour levels [0, 50, 100, 150, ....]\n",
    "levels = [5 + 5*n for n in range(15)]\n",
    "\n",
    "# Make the contour plot\n",
    "a = axes[0].contourf(to_np(wspd_cross), cmap=get_cmap(\"jet\"))\n",
    "# Add the color bar\n",
    "fig.colorbar(a, ax=axes[0])\n",
    "\n",
    "b = axes[1].contourf(to_np(dbz_cross), levels=levels, cmap=get_cmap(\"jet\"))\n",
    "fig.colorbar(b, ax=axes[1])\n",
    "\n",
    "# Set the x-ticks to use latitude and longitude labels.\n",
    "coord_pairs = to_np(dbz_cross.coords[\"xy_loc\"])\n",
    "x_ticks = np.arange(coord_pairs.shape[0])\n",
    "x_labels = [pair.latlon_str() for pair in to_np(coord_pairs)]\n",
    "axes[0].set_xticks(x_ticks[::20])\n",
    "axes[0].set_xticklabels([], rotation=45)\n",
    "axes[1].set_xticks(x_ticks[::20])\n",
    "axes[1].set_xticklabels(x_labels[::20], rotation=45, fontsize=6) \n",
    "\n",
    "\n",
    "# Set the y-ticks to be height.\n",
    "vert_vals = to_np(dbz_cross.coords[\"vertical\"])\n",
    "v_ticks = np.arange(vert_vals.shape[0])\n",
    "axes[0].set_yticks(v_ticks[::20])\n",
    "axes[0].set_yticklabels(vert_vals[::20], fontsize=6) \n",
    "axes[1].set_yticks(v_ticks[::20])\n",
    "axes[1].set_yticklabels(vert_vals[::20], fontsize=6) \n",
    "\n",
    "# Set the x-axis and  y-axis labels\n",
    "axes[1].set_xlabel(\"Latitude, Longitude\", fontsize=7)\n",
    "axes[0].set_ylabel(\"Height (m)\", fontsize=7)\n",
    "axes[1].set_ylabel(\"Height (m)\", fontsize=7)\n",
    "\n",
    "# Add a title\n",
    "axes[0].set_title(\"Cross-Section of Wind Speed (kt)\", {\"fontsize\" : 10})\n",
    "axes[1].set_title(\"Cross-Section of Reflectivity (dBZ)\", {\"fontsize\" : 10})\n",
    "\n",
    "plt.savefig(\"/Users/ladwig/Documents/workspace/wrf_python/doc/source/_static/images/matthew_cross.png\",\n",
    "           transparent=True, bbox_inches=\"tight\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Cartopy Panel Plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from __future__ import (absolute_import, division, print_function, unicode_literals)\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.cm import get_cmap\n",
    "import cartopy.crs as crs\n",
    "import cartopy.feature as cfeature\n",
    "from netCDF4 import Dataset\n",
    "\n",
    "from wrf import (getvar, to_np, vertcross, smooth2d, CoordPair, GeoBounds, get_cartopy, \n",
    "                 latlon_coords, cartopy_xlim, cartopy_ylim)\n",
    "\n",
    "# Open the NetCDF file\n",
    "ncfile = Dataset(\"/Users/ladwig/Documents/wrf_files/wrfout_d01_2016-10-07_00_00_00\")\n",
    "\n",
    "# Get the WRF variables\n",
    "slp = getvar(ncfile, \"slp\")\n",
    "smooth_slp = smooth2d(slp, 3)\n",
    "ctt = getvar(ncfile, \"ctt\")\n",
    "z = getvar(ncfile, \"z\", timeidx=0)\n",
    "dbz = getvar(ncfile, \"dbz\", timeidx=0)\n",
    "Z = 10**(dbz/10.)\n",
    "wspd =  getvar(ncfile, \"wspd_wdir\", units=\"kt\")[0,:]\n",
    "\n",
    "# Set the start point and end point for the cross section\n",
    "start_point = CoordPair(lat=26.76, lon=-80.0)\n",
    "end_point = CoordPair(lat=26.76, lon=-77.8)\n",
    "\n",
    "# Compute the vertical cross-section interpolation.  Also, include the lat/lon points along the cross-section \n",
    "# in the metadata by setting latlon to True.\n",
    "z_cross = vertcross(Z, z, wrfin=ncfile, start_point=start_point, end_point=end_point, latlon=True, meta=True)\n",
    "wspd_cross = vertcross(wspd, z, wrfin=ncfile, start_point=start_point, end_point=end_point, latlon=True, meta=True)\n",
    "dbz_cross = 10.0 * np.log10(z_cross)\n",
    "\n",
    "# Get the lat/lon points\n",
    "lats, lons = latlon_coords(slp)\n",
    "\n",
    "# Get the cartopy projection object\n",
    "cart_proj = get_cartopy(slp)\n",
    "\n",
    "# Create a figure that will have 3 subplots\n",
    "fig = plt.figure(figsize=(12,9))\n",
    "ax_ctt = fig.add_subplot(1,2,1,projection=cart_proj)\n",
    "ax_wspd = fig.add_subplot(2,2,2)\n",
    "ax_dbz = fig.add_subplot(2,2,4)\n",
    "\n",
    "# Download and create the states, land, and oceans using cartopy features\n",
    "states = cfeature.NaturalEarthFeature(category='cultural', scale='50m', facecolor='none',\n",
    "                             name='admin_1_states_provinces_shp')\n",
    "land = cfeature.NaturalEarthFeature(category='physical', name='land', scale='50m', \n",
    "                                    facecolor=cfeature.COLORS['land'])\n",
    "ocean = cfeature.NaturalEarthFeature(category='physical', name='ocean', scale='50m', \n",
    "                                     facecolor=cfeature.COLORS['water'])\n",
    "\n",
    "# Make the pressure contours\n",
    "contour_levels = [960, 965, 970, 975, 980, 990]\n",
    "c1 = ax_ctt.contour(lons, lats, to_np(smooth_slp), levels=contour_levels, colors=\"white\", \n",
    "            transform=crs.PlateCarree(), zorder=3, linewidths=1.0)\n",
    "\n",
    "# Create the filled cloud top temperature contours\n",
    "contour_levels = [-80.0, -70.0, -60, -50, -40, -30, -20, -10, 0, 10]\n",
    "ctt_contours = ax_ctt.contourf(to_np(lons), to_np(lats), to_np(ctt), contour_levels, cmap=get_cmap(\"Greys\"),\n",
    "             transform=crs.PlateCarree(), zorder=2)\n",
    "\n",
    "ax_ctt.plot([start_point.lon, end_point.lon], [start_point.lat, end_point.lat], color=\"yellow\", \n",
    "            marker=\"o\", transform=crs.PlateCarree(), zorder=3)\n",
    "\n",
    "# Create the color bar for cloud top temperature\n",
    "cb_ctt = fig.colorbar(ctt_contours, ax=ax_ctt, shrink=.60)\n",
    "cb_ctt.ax.tick_params(labelsize=5)\n",
    "\n",
    "# Draw the oceans, land, and states\n",
    "ax_ctt.add_feature(ocean)\n",
    "ax_ctt.add_feature(land)\n",
    "ax_ctt.add_feature(states, linewidth=.5, edgecolor=\"black\")\n",
    "\n",
    "# Crop the domain to the region around the hurricane\n",
    "hur_bounds = GeoBounds(CoordPair(lat=np.amin(to_np(lats)), lon=-85.0),\n",
    "                       CoordPair(lat=30.0, lon=-72.0))\n",
    "ax_ctt.set_xlim(cartopy_xlim(ctt, geobounds=hur_bounds))\n",
    "ax_ctt.set_ylim(cartopy_ylim(ctt, geobounds=hur_bounds))\n",
    "ax_ctt.gridlines(color=\"white\", linestyle=\"dotted\")\n",
    "\n",
    "# Make the contour plot for wind speed\n",
    "wspd_contours = ax_wspd.contourf(to_np(wspd_cross), cmap=get_cmap(\"jet\"))\n",
    "# Add the color bar\n",
    "cb_wspd = fig.colorbar(wspd_contours, ax=ax_wspd)\n",
    "cb_wspd.ax.tick_params(labelsize=6)\n",
    "\n",
    "# Make the contour plot for dbz\n",
    "levels = [5 + 5*n for n in range(15)]\n",
    "dbz_contours = ax_dbz.contourf(to_np(dbz_cross), levels=levels, cmap=get_cmap(\"jet\"))\n",
    "cb_dbz = fig.colorbar(dbz_contours, ax=ax_dbz)\n",
    "cb_dbz.ax.tick_params(labelsize=6)\n",
    "\n",
    "# Set the x-ticks to use latitude and longitude labels\n",
    "coord_pairs = to_np(dbz_cross.coords[\"xy_loc\"])\n",
    "x_ticks = np.arange(coord_pairs.shape[0])\n",
    "x_labels = [pair.latlon_str() for pair in to_np(coord_pairs)]\n",
    "ax_wspd.set_xticks(x_ticks[::20])\n",
    "ax_wspd.set_xticklabels([], rotation=45)\n",
    "ax_dbz.set_xticks(x_ticks[::20])\n",
    "ax_dbz.set_xticklabels(x_labels[::20], rotation=45, fontsize=6) \n",
    "\n",
    "# Set the y-ticks to be height\n",
    "vert_vals = to_np(dbz_cross.coords[\"vertical\"])\n",
    "v_ticks = np.arange(vert_vals.shape[0])\n",
    "ax_wspd.set_yticks(v_ticks[::20])\n",
    "ax_wspd.set_yticklabels(vert_vals[::20], fontsize=6) \n",
    "ax_dbz.set_yticks(v_ticks[::20])\n",
    "ax_dbz.set_yticklabels(vert_vals[::20], fontsize=6) \n",
    "\n",
    "# Set the x-axis and  y-axis labels\n",
    "ax_dbz.set_xlabel(\"Latitude, Longitude\", fontsize=8)\n",
    "ax_wspd.set_ylabel(\"Height (m)\", fontsize=8)\n",
    "ax_dbz.set_ylabel(\"Height (m)\", fontsize=8)\n",
    "\n",
    "# Add a title\n",
    "ax_ctt.set_title(\"Cloud Top Temperature (degC)\", {\"fontsize\" : 10})\n",
    "ax_wspd.set_title(\"Cross-Section of Wind Speed (kt)\", {\"fontsize\" : 10})\n",
    "ax_dbz.set_title(\"Cross-Section of Reflectivity (dBZ)\", {\"fontsize\" : 10})\n",
    "\n",
    "plt.savefig(\"/Users/ladwig/Documents/workspace/wrf_python/doc/source/_static/images/matthew.png\",\n",
    "           transparent=True, bbox_inches=\"tight\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from __future__ import (absolute_import, division, print_function, unicode_literals)\n",
    "    \n",
    "from netCDF4 import Dataset   \n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.cm import get_cmap\n",
    "import cartopy.crs as crs\n",
    "from cartopy.feature import NaturalEarthFeature\n",
    "\n",
    "from wrf import to_np, getvar, smooth2d, get_cartopy, cartopy_xlim, cartopy_ylim\n",
    "\n",
    "# Open the NetCDF file\n",
    "ncfile = Dataset(\"/Users/ladwig/Documents/wrf_files/wrfout_d01_2016-10-07_00_00_00\")\n",
    "\n",
    "# Get the sea level pressure\n",
    "slp = getvar(ncfile, \"slp\")\n",
    "\n",
    "# Smooth the sea level pressure since it tends to be noisy near the mountains\n",
    "smooth_slp = smooth2d(slp, 3, cenweight=4)\n",
    "\n",
    "# Get the latitude and longitude points\n",
    "lats, lons = latlon_coords(slp)\n",
    "\n",
    "# Get the cartopy mapping object\n",
    "cart_proj = get_cartopy(slp)\n",
    "\n",
    "# Create a figure\n",
    "fig = plt.figure(figsize=(12,6.0))\n",
    "# Set the GeoAxes to the projection used by WRF\n",
    "ax = plt.axes(projection=cart_proj)\n",
    "\n",
    "# Download and add the states and coastlines\n",
    "states = NaturalEarthFeature(category='cultural', scale='50m', facecolor='none',\n",
    "                             name='admin_1_states_provinces_shp')\n",
    "ax.add_feature(states, linewidth=.5, edgecolor=\"black\")\n",
    "ax.coastlines('50m', linewidth=0.8, zorder=3)\n",
    "\n",
    "# Make the contour outlines and filled contours for the smoothed sea level pressure.\n",
    "plt.contour(to_np(lons), to_np(lats), to_np(smooth_slp), 10, colors=\"black\", transform=crs.PlateCarree())\n",
    "plt.contourf(to_np(lons), to_np(lats), to_np(smooth_slp), 10, transform=crs.PlateCarree(), cmap=get_cmap(\"jet\"))\n",
    "\n",
    "# Add a color bar\n",
    "plt.colorbar(ax=ax, shrink=.98)\n",
    "\n",
    "# Set the map limits.  Not really necessary, but used for demonstration.\n",
    "ax.set_xlim(cartopy_xlim(smooth_slp))\n",
    "ax.set_ylim(cartopy_ylim(smooth_slp))\n",
    "\n",
    "# Add the gridlines\n",
    "ax.gridlines(color=\"black\", linestyle=\"dotted\")\n",
    "\n",
    "plt.title(\"Sea Level Pressure (hPa)\")\n",
    "\n",
    "plt.savefig(\"/Users/ladwig/Documents/workspace/wrf_python/doc/source/_static/images/cartopy_slp.png\",\n",
    "           transparent=True, bbox_inches=\"tight\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Vertical Cross Section"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import (absolute_import, division, print_function, unicode_literals)\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.cm import get_cmap\n",
    "import cartopy.crs as crs\n",
    "from cartopy.feature import NaturalEarthFeature\n",
    "from netCDF4 import Dataset\n",
    "\n",
    "from wrf import to_np, getvar, CoordPair, vertcross\n",
    "\n",
    "# Open the NetCDF file\n",
    "filename = \"/Users/ladwig/Documents/wrf_files/wrfout_d01_2016-10-07_00_00_00\"\n",
    "ncfile = Dataset(filename)\n",
    "\n",
    "# Extract the model height and wind speed\n",
    "z = getvar(ncfile, \"z\")\n",
    "wspd =  getvar(ncfile, \"uvmet_wspd_wdir\", units=\"kt\")[0,:]\n",
    "\n",
    "# Create the start point and end point for the cross section\n",
    "start_point = CoordPair(lat=26.76, lon=-80.0)\n",
    "end_point = CoordPair(lat=26.76, lon=-77.8)\n",
    "\n",
    "# Compute the vertical cross-section interpolation.  Also, include the lat/lon points along the cross-section.\n",
    "wspd_cross = vertcross(wspd, z, wrfin=ncfile, start_point=start_point, end_point=end_point, latlon=True, meta=True)\n",
    "\n",
    "# Create the figure\n",
    "fig = plt.figure(figsize=(12,6))\n",
    "ax = plt.axes()\n",
    "\n",
    "# Make the contour plot\n",
    "wspd_contours = ax.contourf(to_np(wspd_cross), cmap=get_cmap(\"jet\"))\n",
    "\n",
    "# Add the color bar\n",
    "plt.colorbar(wspd_contours, ax=ax)\n",
    "\n",
    "# Set the x-ticks to use latitude and longitude labels.\n",
    "coord_pairs = to_np(wspd_cross.coords[\"xy_loc\"])\n",
    "x_ticks = np.arange(coord_pairs.shape[0])\n",
    "x_labels = [pair.latlon_str(fmt=\"{:.2f}, {:.2f}\") for pair in to_np(coord_pairs)]\n",
    "ax.set_xticks(x_ticks[::20])\n",
    "ax.set_xticklabels(x_labels[::20], rotation=45, fontsize=8) \n",
    "\n",
    "# Set the y-ticks to be height.\n",
    "vert_vals = to_np(wspd_cross.coords[\"vertical\"])\n",
    "v_ticks = np.arange(vert_vals.shape[0])\n",
    "ax.set_yticks(v_ticks[::20])\n",
    "ax.set_yticklabels(vert_vals[::20], fontsize=8) \n",
    "\n",
    "# Set the x-axis and  y-axis labels\n",
    "ax.set_xlabel(\"Latitude, Longitude\", fontsize=12)\n",
    "ax.set_ylabel(\"Height (m)\", fontsize=12)\n",
    "\n",
    "plt.title(\"Vertical Cross Section of Wind Speed (kt)\")\n",
    "\n",
    "plt.savefig(\"/Users/ladwig/Documents/workspace/wrf_python/doc/source/_static/images/cartopy_cross.png\",\n",
    "           transparent=True, bbox_inches=\"tight\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Vertical Cross Section with Mountains"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from matplotlib import pyplot\n",
    "from matplotlib.cm import get_cmap\n",
    "from matplotlib.colors import from_levels_and_colors\n",
    "from cartopy import crs\n",
    "from cartopy.feature import NaturalEarthFeature, COLORS\n",
    "from netCDF4 import Dataset\n",
    "from wrf import (getvar, to_np, get_cartopy, latlon_coords, vertcross,\n",
    "                 cartopy_xlim, cartopy_ylim, interpline, CoordPair)\n",
    "\n",
    "wrf_file = Dataset(\"/Users/ladwig/Documents/wrf_files/boise_tutorial/orig/wrfout_d01_2010-06-04_00:00:00\")\n",
    "\n",
    "# Define the cross section start and end points\n",
    "cross_start = CoordPair(lat=43.5, lon=-116.5)\n",
    "cross_end = CoordPair(lat=43.5, lon=-114)\n",
    "\n",
    "# Get the WRF variables\n",
    "ht = getvar(wrf_file, \"z\", timeidx=-1)\n",
    "ter = getvar(wrf_file, \"ter\", timeidx=-1)\n",
    "dbz = getvar(wrf_file, \"dbz\", timeidx=-1)\n",
    "max_dbz = getvar(wrf_file, \"mdbz\", timeidx=-1)\n",
    "Z = 10**(dbz/10.) # Use linear Z for interpolation\n",
    "\n",
    "# Compute the vertical cross-section interpolation.  Also, include the lat/lon\n",
    "# points along the cross-section in the metadata by setting latlon to True.\n",
    "z_cross = vertcross(Z, ht, wrfin=wrf_file, \n",
    "                    start_point=cross_start, \n",
    "                    end_point=cross_end,\n",
    "                    latlon=True, meta=True)\n",
    "\n",
    "# Convert back to dBz after interpolation\n",
    "dbz_cross = 10.0 * np.log10(z_cross)\n",
    "\n",
    "# Add back the attributes that xarray dropped from the operations above\n",
    "dbz_cross.attrs.update(z_cross.attrs)\n",
    "dbz_cross.attrs[\"description\"] = \"radar reflectivity cross section\"\n",
    "dbz_cross.attrs[\"units\"] = \"dBZ\"\n",
    "\n",
    "# To remove the slight gap between the dbz and terrain due to contouring, the new vertical \n",
    "# grid spacing, and grid staggering, let's fill in the lower grid cells with the first non-missing \n",
    "# value for each column.\n",
    "\n",
    "# Make a copy of the z cross data. Let's use regular numpy arrays for this.\n",
    "dbz_cross_filled = np.ma.copy(to_np(dbz_cross))\n",
    "\n",
    "# For each cross section column, find the first index with non-missing values and copy\n",
    "# these to the missing elements below.\n",
    "for i in range(dbz_cross_filled.shape[-1]):\n",
    "    column_vals = dbz_cross_filled[:,i]\n",
    "    # Let's find the lowest index that isn't filled. The nonzero function \n",
    "    # finds all unmasked values greater than 0. Since 0 is a valid value\n",
    "    # for dBZ, let's change that threshold to be -200 dBZ instead. \n",
    "    first_idx = int(np.transpose((column_vals > -200).nonzero())[0])\n",
    "    dbz_cross_filled[0:first_idx, i] = dbz_cross_filled[first_idx, i]\n",
    "            \n",
    "# Get the terrain heights along the cross section line\n",
    "ter_line = interpline(ter, wrfin=wrf_file, start_point=cross_start, end_point=cross_end)\n",
    "\n",
    "# Get the lat/lon points\n",
    "lats, lons = latlon_coords(dbz)\n",
    "\n",
    "# Get the cartopy projection object\n",
    "cart_proj = get_cartopy(dbz)\n",
    "\n",
    "# Create a figure that will have 2 subplots (1 row, 2 columns)\n",
    "fig = pyplot.figure(figsize=(12,9))\n",
    "ax_cross = pyplot.axes()\n",
    "\n",
    "dbz_levels = np.arange(5., 75., 5.)\n",
    "\n",
    "# This is the NWS color table.\n",
    "dbz_rgb = np.array([[4,233,231],\n",
    "                    [1,159,244], [3,0,244],\n",
    "                    [2,253,2], [1,197,1],\n",
    "                    [0,142,0], [253,248,2],\n",
    "                    [229,188,0], [253,149,0],\n",
    "                    [253,0,0], [212,0,0],\n",
    "                    [188,0,0],[248,0,253],\n",
    "                    [152,84,198]], np.float32) / 255.0\n",
    "    \n",
    "dbz_map, dbz_norm = from_levels_and_colors(dbz_levels, dbz_rgb, extend=\"max\")\n",
    "\n",
    "# Make the cross section plot for dbz\n",
    "dbz_levels = np.arange(5.,75.,5.)\n",
    "xs = np.arange(0, dbz_cross.shape[-1], 1)\n",
    "ys = to_np(dbz_cross.coords[\"vertical\"])\n",
    "dbz_contours = ax_cross.contourf(xs, \n",
    "                                 ys, \n",
    "                                 to_np(dbz_cross_filled), \n",
    "                                 levels=dbz_levels,\n",
    "                                 cmap=dbz_map, \n",
    "                                 norm=dbz_norm, \n",
    "                                 extend=\"max\")\n",
    "cb_dbz = fig.colorbar(dbz_contours, ax=ax_cross)\n",
    "cb_dbz.ax.tick_params(labelsize=8)\n",
    "\n",
    "# Fill in the mountain area\n",
    "ht_fill = ax_cross.fill_between(xs, 0, to_np(ter_line), facecolor=\"saddlebrown\")\n",
    "\n",
    "# Set the x-ticks to use latitude and longitude labels\n",
    "coord_pairs = to_np(dbz_cross.coords[\"xy_loc\"])\n",
    "x_ticks = np.arange(coord_pairs.shape[0])\n",
    "x_labels = [pair.latlon_str() for pair in to_np(coord_pairs)]\n",
    "\n",
    "# Set the desired number of x ticks below\n",
    "num_ticks = 5\n",
    "thin = int((len(x_ticks) / num_ticks) + .5)\n",
    "ax_cross.set_xticks(x_ticks[::thin])\n",
    "ax_cross.set_xticklabels(x_labels[::thin], rotation=45, fontsize=8)\n",
    "\n",
    "# Set the x-axis and  y-axis labels\n",
    "ax_cross.set_xlabel(\"Latitude, Longitude\", fontsize=12)\n",
    "ax_cross.set_ylabel(\"Height (m)\", fontsize=12)\n",
    "\n",
    "# Add a title\n",
    "ax_cross.set_title(\"Cross-Section of Reflectivity (dBZ)\", {\"fontsize\" : 14})\n",
    "\n",
    "plt.savefig(\"/Users/ladwig/Documents/workspace/wrf_python/doc/source/_static/images/cross_mtns.png\",\n",
    "           transparent=True, bbox_inches=\"tight\")\n",
    "\n",
    "pyplot.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Wind Barbs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import (absolute_import, division, print_function, unicode_literals)\n",
    "\n",
    "from netCDF4 import Dataset \n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.cm import get_cmap\n",
    "import cartopy.crs as crs\n",
    "from cartopy.feature import NaturalEarthFeature\n",
    "\n",
    "from wrf import getvar, interplevel, to_np, latlon_coords, get_cartopy, cartopy_xlim, cartopy_ylim\n",
    "\n",
    "# Open the NetCDF file\n",
    "ncfile = Dataset(\"/Users/ladwig/Documents/wrf_files/wrfout_d01_2016-10-07_00_00_00\")\n",
    "\n",
    "# Extract the pressure, geopotential height, and wind variables\n",
    "p = getvar(ncfile, \"pressure\")\n",
    "z = getvar(ncfile, \"z\", units=\"dm\")\n",
    "ua = getvar(ncfile, \"ua\", units=\"kt\")\n",
    "va = getvar(ncfile, \"va\", units=\"kt\")\n",
    "wspd = getvar(ncfile, \"wspd_wdir\", units=\"kts\")[0,:]\n",
    "\n",
    "# Interpolate geopotential height, u, and v winds to 500 hPa \n",
    "ht_500 = interplevel(z, p, 500)\n",
    "u_500 = interplevel(ua, p, 500)\n",
    "v_500 = interplevel(va, p, 500)\n",
    "wspd_500 = interplevel(wspd, p, 500)\n",
    "\n",
    "# Get the lat/lon coordinates\n",
    "lats, lons = latlon_coords(ht_500)\n",
    "\n",
    "# Get the map projection information\n",
    "cart_proj = get_cartopy(ht_500)\n",
    "\n",
    "# Create the figure\n",
    "fig = plt.figure(figsize=(12,9))\n",
    "ax = plt.axes(projection=cart_proj)\n",
    "\n",
    "# Download and add the states and coastlines\n",
    "states = NaturalEarthFeature(category='cultural', scale='50m', facecolor='none',\n",
    "                             name='admin_1_states_provinces_shp')\n",
    "ax.add_feature(states, linewidth=0.5, edgecolor=\"black\")\n",
    "ax.coastlines('50m', linewidth=0.8)\n",
    "\n",
    "# Add the 500 hPa geopotential height contours\n",
    "levels = np.arange(520., 580., 6.)\n",
    "contours = plt.contour(to_np(lons), to_np(lats), to_np(ht_500), levels=levels, colors=\"black\", \n",
    "            transform=crs.PlateCarree())\n",
    "plt.clabel(contours, inline=1, fontsize=10, fmt=\"%i\")\n",
    "\n",
    "# Add the wind speed contours\n",
    "levels = [25, 30, 35, 40, 50, 60, 70, 80, 90, 100, 110, 120]\n",
    "wspd_contours = plt.contourf(to_np(lons), to_np(lats), to_np(wspd_500), levels=levels,\n",
    "                             cmap=get_cmap(\"rainbow\"), \n",
    "                             transform=crs.PlateCarree())\n",
    "plt.colorbar(wspd_contours, ax=ax, orientation=\"horizontal\", pad=.05)\n",
    "\n",
    "# Add the 500 hPa wind barbs, only plotting every 125th data point.\n",
    "plt.barbs(to_np(lons[::125,::125]), to_np(lats[::125,::125]), to_np(u_500[::125, ::125]), \n",
    "          to_np(v_500[::125, ::125]), transform=crs.PlateCarree(), length=6)\n",
    "\n",
    "# Set the map bounds\n",
    "ax.set_xlim(cartopy_xlim(ht_500))\n",
    "ax.set_ylim(cartopy_ylim(ht_500))\n",
    "\n",
    "ax.gridlines()\n",
    "\n",
    "plt.title(\"500 MB Height (dm), Wind Speed (kt), Barbs (kt)\")\n",
    "\n",
    "plt.savefig(\"/Users/ladwig/Documents/workspace/wrf_python/doc/source/_static/images/cartopy_500.png\",\n",
    "           transparent=True, bbox_inches=\"tight\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "    \n",
    "from netCDF4 import Dataset\n",
    "from wrf import getvar, get_cartopy, latlon_coords, geo_bounds\n",
    "\n",
    "ncfile = Dataset(\"/Users/ladwig/Documents/wrf_files/wrfout_d01_2016-10-07_00_00_00\")\n",
    "\n",
    "slp = getvar(ncfile, \"slp\")\n",
    "\n",
    "# Get the cartopy mapping object\n",
    "cart_proj = get_cartopy(slp)\n",
    "\n",
    "print (cart_proj)\n",
    "\n",
    "# Get the latitude and longitude coordinate.  This is needed for plotting.\n",
    "lats, lons = latlon_coords(slp)\n",
    "\n",
    "# Get the geobounds for the full domain\n",
    "bounds = geo_bounds(slp)\n",
    "print (bounds)\n",
    "\n",
    "# Get the geographic boundaries for a subset of the domain\n",
    "slp_subset = slp[150:250, 150:250]\n",
    "slp_subset_bounds = geo_bounds(slp_subset)\n",
    "\n",
    "print (slp_subset_bounds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "    \n",
    "from netCDF4 import Dataset\n",
    "from wrf import get_cartopy, geo_bounds\n",
    "\n",
    "ncfile = Dataset(\"/Users/ladwig/Documents/wrf_files/wrfout_d01_2016-10-07_00_00_00\")\n",
    "\n",
    "cart_proj = get_cartopy(wrfin=ncfile)\n",
    "\n",
    "print (cart_proj)\n",
    "\n",
    "bounds = geo_bounds(wrfin=ncfile)\n",
    "\n",
    "print (bounds)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Basemap Examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "from __future__ import (absolute_import, division, print_function, unicode_literals)\n",
    "    \n",
    "from netCDF4 import Dataset   \n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.cm import get_cmap\n",
    "from mpl_toolkits.basemap import Basemap\n",
    "\n",
    "from wrf import to_np, getvar, smooth2d, get_basemap, latlon_coords\n",
    "\n",
    "# Open the NetCDF file\n",
    "ncfile = Dataset(\"/Users/ladwig/Documents/wrf_files/wrfout_d01_2016-10-07_00_00_00\")\n",
    "\n",
    "# Get the sea level pressure\n",
    "slp = getvar(ncfile, \"slp\")\n",
    "\n",
    "# Smooth the sea level pressure since it tends to be noisy near the mountains\n",
    "smooth_slp = smooth2d(slp, 3, cenweight=4)\n",
    "\n",
    "# Get the latitude and longitude points\n",
    "lats, lons = latlon_coords(slp)\n",
    "\n",
    "# Get the basemap object\n",
    "bm = get_basemap(slp)\n",
    "\n",
    "# Create a figure\n",
    "fig = plt.figure(figsize=(12,6))\n",
    "\n",
    "# Add geographic outlines\n",
    "bm.drawcoastlines(linewidth=0.25)\n",
    "bm.drawstates(linewidth=0.25)\n",
    "bm.drawcountries(linewidth=0.25)\n",
    "\n",
    "# Convert the lats and lons to x and y.  Make sure you convert the lats and lons to \n",
    "# numpy arrays via to_np, or basemap crashes with an undefined RuntimeError.\n",
    "x, y = bm(to_np(lons), to_np(lats))\n",
    "\n",
    "# Draw the contours and filled contours\n",
    "bm.contour(x, y, to_np(smooth_slp), 10, colors=\"black\")\n",
    "bm.contourf(x, y, to_np(smooth_slp), 10, cmap=get_cmap(\"jet\"))\n",
    "\n",
    "# Add a color bar\n",
    "plt.colorbar(shrink=.98)\n",
    "\n",
    "plt.title(\"Sea Level Pressure (hPa)\")\n",
    "\n",
    "plt.savefig(\"/Users/ladwig/Documents/workspace/wrf_python/doc/source/_static/images/basemap_slp.png\",\n",
    "           transparent=True, bbox_inches=\"tight\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import (absolute_import, division, print_function, unicode_literals)\n",
    "\n",
    "from netCDF4 import Dataset \n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.cm import get_cmap\n",
    "\n",
    "from wrf import getvar, interplevel, to_np, get_basemap, latlon_coords\n",
    "\n",
    "# Open the NetCDF file\n",
    "ncfile = Dataset(\"/Users/ladwig/Documents/wrf_files/wrfout_d01_2016-10-07_00_00_00\")\n",
    "\n",
    "# Extract the pressure, geopotential height, and wind variables\n",
    "p = getvar(ncfile, \"pressure\")\n",
    "z = getvar(ncfile, \"z\", units=\"dm\")\n",
    "ua = getvar(ncfile, \"ua\", units=\"kt\")\n",
    "va = getvar(ncfile, \"va\", units=\"kt\")\n",
    "wspd = getvar(ncfile, \"wspd_wdir\", units=\"kts\")[0,:]\n",
    "\n",
    "# Interpolate geopotential height, u, and v winds to 500 hPa \n",
    "ht_500 = interplevel(z, p, 500)\n",
    "u_500 = interplevel(ua, p, 500)\n",
    "v_500 = interplevel(va, p, 500)\n",
    "wspd_500 = interplevel(wspd, p, 500)\n",
    "\n",
    "# Get the lat/lon coordinates\n",
    "lats, lons = latlon_coords(ht_500)\n",
    "\n",
    "# Get the basemap object\n",
    "bm = get_basemap(ht_500)\n",
    "\n",
    "# Create the figure\n",
    "fig = plt.figure(figsize=(12,9))\n",
    "ax = plt.axes()\n",
    "\n",
    "# Convert the lat/lon coordinates to x/y coordinates in the projection space\n",
    "x, y = bm(to_np(lons), to_np(lats))\n",
    "\n",
    "# Add the 500 hPa geopotential height contours\n",
    "levels = np.arange(520., 580., 6.)\n",
    "contours = bm.contour(x, y, to_np(ht_500), levels=levels, colors=\"black\")\n",
    "plt.clabel(contours, inline=1, fontsize=10, fmt=\"%i\")\n",
    "\n",
    "# Add the wind speed contours\n",
    "levels = [25, 30, 35, 40, 50, 60, 70, 80, 90, 100, 110, 120]\n",
    "wspd_contours = bm.contourf(x, y, to_np(wspd_500), levels=levels,\n",
    "                             cmap=get_cmap(\"rainbow\"))\n",
    "plt.colorbar(wspd_contours, ax=ax, orientation=\"horizontal\", pad=.05)\n",
    "\n",
    "# Add the geographic boundaries\n",
    "bm.drawcoastlines(linewidth=0.25)\n",
    "bm.drawstates(linewidth=0.25)\n",
    "bm.drawcountries(linewidth=0.25)\n",
    "\n",
    "# Add the 500 hPa wind barbs, only plotting every 125th data point.\n",
    "bm.barbs(x[::125,::125], y[::125,::125], to_np(u_500[::125, ::125]), \n",
    "          to_np(v_500[::125, ::125]), length=6)\n",
    "\n",
    "plt.title(\"500 MB Height (dm), Wind Speed (kt), Barbs (kt)\")\n",
    "\n",
    "plt.savefig(\"/Users/ladwig/Documents/workspace/wrf_python/doc/source/_static/images/basemap_500.png\",\n",
    "           transparent=True, bbox_inches=\"tight\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Basemap Panel Plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.cm import get_cmap\n",
    "from netCDF4 import Dataset\n",
    "\n",
    "from wrf import getvar, to_np, vertcross, smooth2d, CoordPair, get_basemap, latlon_coords\n",
    "\n",
    "# Open the NetCDF file\n",
    "ncfile = Dataset(\"/Users/ladwig/Documents/wrf_files/wrfout_d01_2016-10-07_00_00_00\")\n",
    "\n",
    "# Get the WRF variables\n",
    "slp = getvar(ncfile, \"slp\")\n",
    "smooth_slp = smooth2d(slp, 3)\n",
    "ctt = getvar(ncfile, \"ctt\")\n",
    "z = getvar(ncfile, \"z\", timeidx=0)\n",
    "dbz = getvar(ncfile, \"dbz\", timeidx=0)\n",
    "Z = 10**(dbz/10.)\n",
    "wspd =  getvar(ncfile, \"wspd_wdir\", units=\"kt\")[0,:]\n",
    "\n",
    "# Set the start point and end point for the cross section\n",
    "start_point = CoordPair(lat=26.76, lon=-80.0)\n",
    "end_point = CoordPair(lat=26.76, lon=-77.8)\n",
    "\n",
    "# Compute the vertical cross-section interpolation.  Also, include the lat/lon points along the cross-section in \n",
    "# the metadata by setting latlon to True.\n",
    "z_cross = vertcross(Z, z, wrfin=ncfile, start_point=start_point, end_point=end_point, latlon=True, meta=True)\n",
    "wspd_cross = vertcross(wspd, z, wrfin=ncfile, start_point=start_point, end_point=end_point, latlon=True, meta=True)\n",
    "dbz_cross = 10.0 * np.log10(z_cross)\n",
    "\n",
    "# Get the latitude and longitude points\n",
    "lats, lons = latlon_coords(slp)\n",
    "\n",
    "# Create the figure that will have 3 subplots\n",
    "fig = plt.figure(figsize=(12,9))\n",
    "ax_ctt = fig.add_subplot(1,2,1)\n",
    "ax_wspd = fig.add_subplot(2,2,2)\n",
    "ax_dbz = fig.add_subplot(2,2,4)\n",
    "\n",
    "# Get the basemap object\n",
    "bm = get_basemap(slp)\n",
    "\n",
    "# Convert the lat/lon points in to x/y points in the projection space\n",
    "x, y = bm(to_np(lons), to_np(lats))\n",
    "\n",
    "# Make the pressure contours\n",
    "contour_levels = [960, 965, 970, 975, 980, 990]\n",
    "c1 = bm.contour(x, y, to_np(smooth_slp), levels=contour_levels, colors=\"white\", \n",
    "                zorder=3, linewidths=1.0, ax=ax_ctt)\n",
    "\n",
    "# Create the filled cloud top temperature contours\n",
    "contour_levels = [-80.0, -70.0, -60, -50, -40, -30, -20, -10, 0, 10]\n",
    "ctt_contours = bm.contourf(x, y, to_np(ctt), contour_levels, cmap=get_cmap(\"Greys\"),\n",
    "                               zorder=2, ax=ax_ctt)\n",
    "\n",
    "point_x, point_y = bm([start_point.lon, end_point.lon], [start_point.lat, end_point.lat])\n",
    "bm.plot([point_x[0], point_x[1]], [point_y[0], point_y[1]], color=\"yellow\", \n",
    "            marker=\"o\", zorder=3, ax=ax_ctt)\n",
    "\n",
    "# Create the color bar for cloud top temperature\n",
    "cb_ctt = fig.colorbar(ctt_contours, ax=ax_ctt, shrink=.60)\n",
    "cb_ctt.ax.tick_params(labelsize=5)\n",
    "\n",
    "# Draw the oceans, land, and states\n",
    "bm.drawcoastlines(linewidth=0.25, ax=ax_ctt)\n",
    "bm.drawstates(linewidth=0.25, ax=ax_ctt)\n",
    "bm.drawcountries(linewidth=0.25, ax=ax_ctt)\n",
    "bm.fillcontinents(color=np.array([ 0.9375 , 0.9375 , 0.859375]), \n",
    "                  ax=ax_ctt, lake_color=np.array([ 0.59375 , 0.71484375, 0.8828125 ]))\n",
    "bm.drawmapboundary(fill_color=np.array([ 0.59375 , 0.71484375, 0.8828125 ]), ax=ax_ctt)\n",
    "\n",
    "# Draw Parallels\n",
    "parallels = np.arange(np.amin(lats), 30., 2.5)\n",
    "bm.drawparallels(parallels, ax=ax_ctt, color=\"white\")\n",
    "\n",
    "merids = np.arange(-85.0, -72.0, 2.5)\n",
    "bm.drawmeridians(merids, ax=ax_ctt, color=\"white\")\n",
    "\n",
    "# Crop the image to the hurricane region\n",
    "x_start, y_start = bm(-85.0, np.amin(lats))\n",
    "x_end, y_end = bm(-72.0, 30.0)\n",
    "\n",
    "ax_ctt.set_xlim([x_start, x_end])\n",
    "ax_ctt.set_ylim([y_start, y_end])\n",
    "\n",
    "# Make the contour plot for wspd\n",
    "wspd_contours = ax_wspd.contourf(to_np(wspd_cross), cmap=get_cmap(\"jet\"))\n",
    "# Add the color bar\n",
    "cb_wspd = fig.colorbar(wspd_contours, ax=ax_wspd)\n",
    "cb_wspd.ax.tick_params(labelsize=6)\n",
    "\n",
    "# Make the contour plot for dbz\n",
    "levels = [5 + 5*n for n in range(15)]\n",
    "dbz_contours = ax_dbz.contourf(to_np(dbz_cross), levels=levels, cmap=get_cmap(\"jet\"))\n",
    "cb_dbz = fig.colorbar(dbz_contours, ax=ax_dbz)\n",
    "cb_dbz.ax.tick_params(labelsize=6)\n",
    "\n",
    "# Set the x-ticks to use latitude and longitude labels.\n",
    "coord_pairs = to_np(dbz_cross.coords[\"xy_loc\"])\n",
    "x_ticks = np.arange(coord_pairs.shape[0])\n",
    "x_labels = [pair.latlon_str() for pair in to_np(coord_pairs)]\n",
    "ax_wspd.set_xticks(x_ticks[::20])\n",
    "ax_wspd.set_xticklabels([], rotation=45)\n",
    "ax_dbz.set_xticks(x_ticks[::20])\n",
    "ax_dbz.set_xticklabels(x_labels[::20], rotation=45, fontsize=6) \n",
    "\n",
    "# Set the y-ticks to be height.\n",
    "vert_vals = to_np(dbz_cross.coords[\"vertical\"])\n",
    "v_ticks = np.arange(vert_vals.shape[0])\n",
    "ax_wspd.set_yticks(v_ticks[::20])\n",
    "ax_wspd.set_yticklabels(vert_vals[::20], fontsize=6) \n",
    "ax_dbz.set_yticks(v_ticks[::20])\n",
    "ax_dbz.set_yticklabels(vert_vals[::20], fontsize=6) \n",
    "\n",
    "# Set the x-axis and  y-axis labels\n",
    "ax_dbz.set_xlabel(\"Latitude, Longitude\", fontsize=8)\n",
    "ax_wspd.set_ylabel(\"Height (m)\", fontsize=8)\n",
    "ax_dbz.set_ylabel(\"Height (m)\", fontsize=8)\n",
    "\n",
    "# Add titles\n",
    "ax_ctt.set_title(\"Cloud Top Temperature (degC)\", {\"fontsize\" : 10})\n",
    "ax_wspd.set_title(\"Cross-Section of Wind Speed (kt)\", {\"fontsize\" : 10})\n",
    "ax_dbz.set_title(\"Cross-Section of Reflectivity (dBZ)\", {\"fontsize\" : 10})\n",
    "\n",
    "plt.savefig(\"/Users/ladwig/Documents/workspace/wrf_python/doc/source/_static/images/basemap_front.png\",\n",
    "           transparent=False, bbox_inches=\"tight\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "    \n",
    "from netCDF4 import Dataset\n",
    "from wrf import getvar, get_basemap, latlon_coords, geo_bounds\n",
    "\n",
    "ncfile = Dataset(\"/Users/ladwig/Documents/wrf_files/wrfout_d01_2016-10-07_00_00_00\")\n",
    "\n",
    "slp = getvar(ncfile, \"slp\")\n",
    "\n",
    "# Get the basemap mapping object\n",
    "basemap_proj = get_basemap(slp)\n",
    "\n",
    "print (basemap_proj)\n",
    "\n",
    "# Get the latitude and longitude coordinate.  This is needed for plotting.\n",
    "lats, lons = latlon_coords(slp)\n",
    "\n",
    "# Get the geobounds for the full domain\n",
    "bounds = geo_bounds(slp)\n",
    "\n",
    "print(bounds)\n",
    "\n",
    "# Get the geographic boundaries for a subset of the domain\n",
    "slp_subset = slp[150:250, 150:250]\n",
    "slp_subset_bounds = geo_bounds(slp_subset)\n",
    "\n",
    "print (slp_subset_bounds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "    \n",
    "from netCDF4 import Dataset\n",
    "from wrf import get_basemap, geo_bounds\n",
    "\n",
    "ncfile = Dataset(\"/Users/ladwig/Documents/wrf_files/wrfout_d01_2016-10-07_00_00_00\")\n",
    "\n",
    "cart_proj = get_basemap(wrfin=ncfile)\n",
    "\n",
    "print (cart_proj)\n",
    "\n",
    "bounds = geo_bounds(wrfin=ncfile)\n",
    "\n",
    "print (bounds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SLP\n",
    "from __future__ import (absolute_import, division, print_function, unicode_literals)\n",
    "    \n",
    "import Ngl\n",
    "import Nio\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "from wrf import to_np, getvar, smooth2d, get_pyngl, latlon_coords\n",
    "\n",
    "ncfile = Nio.open_file(b\"/Users/ladwig/Documents/wrf_files/wrfout_d01_2016-10-07_00_00_00.nc\")\n",
    "\n",
    "# Get the sea level pressure\n",
    "ctt = getvar(ncfile, \"ctt\")\n",
    "\n",
    "wks = Ngl.open_wks(b\"png\", b\"test\")\n",
    "\n",
    "# Get the map projection\n",
    "resources = get_pyngl(ctt)\n",
    "\n",
    "# This needs to be False if you want to crop\n",
    "resources.tfDoNDCOverlay = False\n",
    "\n",
    "resources.mpLeftCornerLatF  = 30.0\n",
    "resources.mpRightCornerLatF = np.amin(to_np(ctt.coords[\"XLAT\"]))\n",
    "resources.mpLeftCornerLonF  = -85.\n",
    "resources.mpRightCornerLonF = -75.\n",
    "\n",
    "lats, lons = latlon_coords(ctt, as_np=True)\n",
    "resources.sfYArray = lats\n",
    "resources.sfXArray = lons\n",
    "\n",
    "resources.cnLevelSelectionMode = b\"ExplicitLevels\" # Define your own\n",
    "resources.cnLevels = np.arange(-80.,30.,10.)\n",
    "resources.cnFillOn = True\n",
    "resources.cnFillMode = b\"RasterFill\"\n",
    "#resources.cnFillPalette = Ngl.read_colormap_file(b\"MPL_Greys\")\n",
    "\n",
    "\n",
    "Ngl.contour_map(wks, to_np(ctt), resources)\n",
    "\n",
    "Ngl.end()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "    \n",
    "from netCDF4 import Dataset\n",
    "from wrf import getvar, get_pyngl, latlon_coords, geo_bounds\n",
    "\n",
    "ncfile = Dataset(\"/Users/ladwig/Documents/wrf_files/wrfout_d01_2016-10-07_00_00_00\")\n",
    "\n",
    "slp = getvar(ncfile, \"slp\")\n",
    "\n",
    "# Get the pyngl mapping object\n",
    "pyngl_resources = get_pyngl(slp)\n",
    "\n",
    "print (pyngl_resources)\n",
    "\n",
    "# Get the latitude and longitude coordinate.  This is needed for plotting.\n",
    "lats, lons = latlon_coords(slp)\n",
    "\n",
    "# Get the geobounds for the full domain\n",
    "bounds = geo_bounds(slp)\n",
    "print(bounds)\n",
    "\n",
    "# Get the geographic boundaries for a subset of the domain\n",
    "slp_subset = slp[150:250, 150:250]\n",
    "slp_subset_bounds = geo_bounds(slp_subset)\n",
    "\n",
    "print (slp_subset_bounds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "    \n",
    "from netCDF4 import Dataset\n",
    "from wrf import get_pyngl, geo_bounds\n",
    "\n",
    "ncfile = Dataset(\"/Users/ladwig/Documents/wrf_files/wrfout_d01_2016-10-07_00_00_00\")\n",
    "\n",
    "pyngl_resources = get_pyngl(wrfin=ncfile)\n",
    "\n",
    "print (pyngl_resources)\n",
    "\n",
    "bounds = geo_bounds(wrfin=ncfile)\n",
    "\n",
    "print (bounds)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# OpenMP Routines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "\n",
    "from wrf import omp_enabled\n",
    "\n",
    "print(omp_enabled())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "\n",
    "from wrf import omp_get_num_procs\n",
    "\n",
    "print(omp_get_num_procs())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "\n",
    "from wrf import omp_set_num_threads, omp_get_max_threads\n",
    "\n",
    "omp_set_num_threads(4)\n",
    "\n",
    "print(omp_get_max_threads())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "\n",
    "from wrf import omp_set_schedule, omp_get_schedule, OMP_SCHED_GUIDED\n",
    "\n",
    "omp_set_schedule(OMP_SCHED_GUIDED, 0)\n",
    "\n",
    "sched, modifier = omp_get_schedule()\n",
    "\n",
    "print(sched, modifier)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Loop and Fill Technique"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function, division\n",
    "\n",
    "import numpy as np\n",
    "from netCDF4 import Dataset\n",
    "from wrf import getvar, ALL_TIMES\n",
    "\n",
    "filename_list = [\"/Users/ladwig/Documents/wrf_files/wrf_vortex_multi/wrfout_d02_2005-08-28_00:00:00\",\n",
    "                 \"/Users/ladwig/Documents/wrf_files/wrf_vortex_multi/wrfout_d02_2005-08-28_12:00:00\", \n",
    "                 \"/Users/ladwig/Documents/wrf_files/wrf_vortex_multi/wrfout_d02_2005-08-29_00:00:00\"]\n",
    "\n",
    "# Result shape (hardcoded for this example, modify as necessary)\n",
    "result_shape = (9, 29, 96, 96)\n",
    "\n",
    "# Only need 4-byte floats\n",
    "z_final = np.empty(result_shape, np.float32)\n",
    "\n",
    "# Modify this number if using more than 1 time per file\n",
    "times_per_file = 4\n",
    "\n",
    "for timeidx in xrange(result_shape[0]):\n",
    "   # Compute the file index and the time index inside the file\n",
    "   fileidx = timeidx // times_per_file\n",
    "   file_timeidx = timeidx % times_per_file\n",
    "\n",
    "   f = Dataset(filename_list[fileidx])  \n",
    "   z = getvar(f, \"z\", file_timeidx)\n",
    "\n",
    "   z_final[timeidx,:] = z[:]\n",
    "   f.close()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using the cache argument"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "\n",
    "import time\n",
    "from netCDF4 import Dataset\n",
    "from wrf import getvar, ALL_TIMES, extract_vars\n",
    "\n",
    "wrf_filenames = [\"/Users/ladwig/Documents/wrf_files/wrf_vortex_multi/wrfout_d02_2005-08-28_00:00:00\",\n",
    "                 \"/Users/ladwig/Documents/wrf_files/wrf_vortex_multi/wrfout_d02_2005-08-28_12:00:00\", \n",
    "                 \"/Users/ladwig/Documents/wrf_files/wrf_vortex_multi/wrfout_d02_2005-08-29_00:00:00\"]\n",
    "\n",
    "wrfin = [Dataset(x) for x in wrf_filenames]\n",
    "\n",
    "start = time.time()\n",
    "my_cache = extract_vars(wrfin, ALL_TIMES, (\"P\", \"PSFC\", \"PB\", \"PH\", \"PHB\", \"T\", \"QVAPOR\", \n",
    "                                           \"HGT\", \"U\", \"V\", \"W\"))\n",
    "end = time.time()\n",
    "print (\"Time taken to build cache: \", (end-start), \"s\")\n",
    "\n",
    "vars = (\"avo\", \"eth\", \"cape_2d\", \"cape_3d\", \"ctt\", \"dbz\", \"mdbz\", \n",
    "        \"geopt\", \"helicity\", \"lat\", \"lon\", \"omg\", \"p\", \"pressure\", \n",
    "        \"pvo\", \"pw\", \"rh2\", \"rh\", \"slp\", \"ter\", \"td2\", \"td\", \"tc\", \n",
    "        \"theta\", \"tk\", \"tv\", \"twb\", \"updraft_helicity\", \"ua\", \"va\", \n",
    "        \"wa\", \"uvmet10\", \"uvmet\", \"z\", \"cfrac\", \"zstag\", \"geopt_stag\")\n",
    "\n",
    "start = time.time()\n",
    "for var in vars:\n",
    "    v = getvar(wrfin, var, ALL_TIMES)\n",
    "end = time.time()\n",
    "no_cache_time = (end-start)\n",
    "\n",
    "print (\"Time taken without variable cache: \", no_cache_time, \"s\")\n",
    "\n",
    "start = time.time()\n",
    "for var in vars:\n",
    "    v = getvar(wrfin, var, ALL_TIMES, cache=my_cache)\n",
    "end = time.time()\n",
    "cache_time = (end-start)\n",
    "\n",
    "print (\"Time taken with variable cache: \", cache_time, \"s\")\n",
    "\n",
    "improvement = ((no_cache_time-cache_time)/no_cache_time) * 100 \n",
    "print (\"The cache decreased computation time by: \", improvement, \"%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using the cache argument with OpenMP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "\n",
    "import time\n",
    "from netCDF4 import Dataset\n",
    "from wrf import getvar, ALL_TIMES, extract_vars, omp_set_num_threads, omp_get_num_procs\n",
    "\n",
    "wrf_filenames = [\"/Users/ladwig/Documents/wrf_files/wrf_vortex_multi/wrfout_d02_2005-08-28_00:00:00\",\n",
    "                 \"/Users/ladwig/Documents/wrf_files/wrf_vortex_multi/wrfout_d02_2005-08-28_12:00:00\", \n",
    "                 \"/Users/ladwig/Documents/wrf_files/wrf_vortex_multi/wrfout_d02_2005-08-29_00:00:00\"]\n",
    "\n",
    "wrfin = [Dataset(x) for x in wrf_filenames]\n",
    "\n",
    "start = time.time()\n",
    "my_cache = extract_vars(wrfin, ALL_TIMES, (\"P\", \"PSFC\", \"PB\", \"PH\", \"PHB\", \"T\", \"QVAPOR\", \n",
    "                                           \"HGT\", \"U\", \"V\", \"W\"))\n",
    "end = time.time()\n",
    "print (\"Time taken to build cache: \", (end-start), \"s\")\n",
    "\n",
    "omp_set_num_threads(omp_get_num_procs())\n",
    "\n",
    "vars = (\"avo\", \"eth\", \"cape_2d\", \"cape_3d\", \"ctt\", \"dbz\", \"mdbz\", \n",
    "        \"geopt\", \"helicity\", \"lat\", \"lon\", \"omg\", \"p\", \"pressure\", \n",
    "        \"pvo\", \"pw\", \"rh2\", \"rh\", \"slp\", \"ter\", \"td2\", \"td\", \"tc\", \n",
    "        \"theta\", \"tk\", \"tv\", \"twb\", \"updraft_helicity\", \"ua\", \"va\", \n",
    "        \"wa\", \"uvmet10\", \"uvmet\", \"z\", \"cfrac\", \"zstag\", \"geopt_stag\")\n",
    "\n",
    "start = time.time()\n",
    "for var in vars:\n",
    "    v = getvar(wrfin, var, ALL_TIMES)\n",
    "end = time.time()\n",
    "no_cache_time = (end-start)\n",
    "\n",
    "print (\"Time taken without variable cache: \", no_cache_time, \"s\")\n",
    "\n",
    "start = time.time()\n",
    "for var in vars:\n",
    "    v = getvar(wrfin, var, ALL_TIMES, cache=my_cache)\n",
    "end = time.time()\n",
    "cache_time = (end-start)\n",
    "\n",
    "print (\"Time taken with variable cache: \", cache_time, \"s\")\n",
    "\n",
    "improvement = ((no_cache_time-cache_time)/no_cache_time) * 100 \n",
    "print (\"The cache decreased computation time by: \", improvement, \"%\")\n",
    "\n",
    "omp_set_num_threads(1)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
